// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <memory>
#include <pollux/expression/complex_view_types.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/expression/vector_readers.h>
#include <pollux/type/floating_point_util.h>
#include <pollux/type/type.h>
#include <pollux/vector/base_vector.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/nulls_builder.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux::functions {
    namespace detail {
        // Below functions return a stock instance of each of the possible errors in
        // SubscriptImpl
        const std::exception_ptr &zeroSubscriptError();

        const std::exception_ptr &badSubscriptError();

        const std::exception_ptr &negativeSubscriptError();

        // A flat vector of map keys, an index into that vector and an index into
        // the original map keys vector that may have encodings.
        struct MapKey {
            const BaseVector *baseVector;
            const vector_size_t baseIndex;
            const vector_size_t index;

            size_t hash() const {
                return baseVector->hash_value_at(baseIndex);
            }

            bool operator==(const MapKey &other) const {
                return baseVector->equal_value_at(
                    other.baseVector, baseIndex, other.baseIndex);
            }

            bool operator<(const MapKey &other) const {
                return baseVector->compare(other.baseVector, baseIndex, other.baseIndex) <
                       0;
            }
        };

        struct MapKeyHasher {
            size_t operator()(const MapKey &key) const {
                return key.hash();
            }
        };

        using MapKeyAllocator = memory::StlAllocator<detail::MapKey>;

        using ComplexKeyHashMap = melon::F14FastSet<
            detail::MapKey,
            detail::MapKeyHasher,
            melon::f14::DefaultKeyEqual<detail::MapKey>,
            MapKeyAllocator>;

        template<typename NativeType>
        class LookupTable;

        class LookupTableBase {
        public:
            template<typename NativeType>
            LookupTable<NativeType> *typedTable() {
                return static_cast<LookupTable<NativeType> *>(this);
            }

            virtual ~LookupTableBase() {
            }
        };

        // NativeType should by TypeTraits<TypeKind>::NativeType for the key's TypeKind.
        template<typename NativeType>
        class LookupTable : public LookupTableBase {
        public:
            LookupTable(memory::MemoryPool &pool)
                : pool_(pool),
                  map_(std::make_unique<outer_map_t>(outer_allocator_t(pool))) {
            }

            auto &map() {
                return map_;
            }

            bool containsMapAtIndex(vector_size_t rowIndex) const {
                return map_->count(rowIndex) != 0;
            }

            void ensureMapAtIndex(vector_size_t rowIndex) const {
                map_->emplace(rowIndex, pool_);
            }

            auto &getMapAtIndex(vector_size_t rowIndex) {
                POLLUX_DCHECK(containsMapAtIndex(rowIndex));
                return map_->find(rowIndex)->second;
            }

        private:
            // If the NativeType is not void, we can materialize the key in memory
            // directly, so we can use a HashMap keyed on the native value.  If it is void
            // then we have to use MapKey as the key to wrap the Vector and avoid
            // materializing the key in memory.
            using inner_allocator_t = std::conditional_t<
                std::is_same_v<NativeType, void>,
                MapKeyAllocator,
                memory::StlAllocator<std::pair<NativeType const, vector_size_t> > >;

            using inner_map_t = std::conditional_t<
                std::is_same_v<NativeType, void>,
                ComplexKeyHashMap,
                typename util::floating_point::HashMapNaNAwareTypeTraits<
                    NativeType,
                    vector_size_t,
                    inner_allocator_t>::Type>;

            using outer_allocator_t =
            memory::StlAllocator<std::pair<vector_size_t const, inner_map_t> >;

            // [rowindex][key] -> offset of value.
            using outer_map_t = melon::F14FastMap<
                vector_size_t,
                inner_map_t,
                melon::f14::DefaultHasher<vector_size_t>,
                melon::f14::DefaultKeyEqual<vector_size_t>,
                outer_allocator_t>;

            memory::MemoryPool &pool_;
            std::unique_ptr<outer_map_t> map_;
        };

        class MapSubscript {
        public:
            explicit MapSubscript(bool allowCaching) : allowCaching_(allowCaching) {
            }

            VectorPtr applyMap(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                exec::EvalCtx &context) const;

            bool cachingEnabled() const {
                return allowCaching_;
            }

            auto &lookupTable() const {
                return lookupTable_;
            }

            auto &firstSeenMap() const {
                return firstSeenMap_;
            }

        private:
            bool shouldTriggerCaching(const VectorPtr &mapArg) const {
                if (!allowCaching_) {
                    return false;
                }

                if (mapArg->type()->childAt(0)->is_boolean()) {
                    // Disable caching if the key type is boolean.
                    allowCaching_ = false;
                    return false;
                }

                if (!firstSeenMap_) {
                    firstSeenMap_ = mapArg;
                    return false;
                }

                if (firstSeenMap_->wrapped_vector() == mapArg->wrapped_vector()) {
                    return true;
                }

                // Disable caching forever.
                allowCaching_ = false;
                lookupTable_.reset();
                firstSeenMap_.reset();
                return false;
            }

            // When true the function is allowed to cache a materialized version of the
            // processed map.
            mutable bool allowCaching_;

            // This is used to check if the same base map is being passed over and over
            // in the function. A shared_ptr is used to guarantee that if the map is
            // seen again then it was not modified.
            mutable VectorPtr firstSeenMap_;

            // Materialized cached version of firstSeenMap_ used to optimize the lookup.
            mutable std::shared_ptr<LookupTableBase> lookupTable_;
        };
    } // namespace detail

    /// Generic subscript/element_at implementation for both array and map data
    /// types.
    ///
    /// Provides four template parameters to configure the behavior:
    /// - allowNegativeIndices: if allowed, negative indices accesses elements
    /// from
    ///   last to the first; otherwise, throws.
    /// - nullOnNegativeIndices: returns NULL for negative indices instead of the
    ///   behavior described above.
    /// - allowOutOfBound: if allowed, returns NULL for out of bound accesses; if
    ///   false, throws an exception.
    /// - indexStartsAtOne: whether indices start at zero or one.
    template<
        bool allowNegativeIndices,
        bool nullOnNegativeIndices,
        bool allowOutOfBound,
        bool indexStartsAtOne>
    class SubscriptImpl : public exec::Subscript {
    public:
        explicit SubscriptImpl(bool allowCaching)
            : mapSubscript_(detail::MapSubscript(allowCaching)) {
        }

        void apply(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args,
            const TypePtr & /* outputType */,
            exec::EvalCtx &context,
            VectorPtr &result) const override {
            POLLUX_CHECK_EQ(args.size(), 2);
            VectorPtr localResult;

            switch (args[0]->type_kind()) {
                case TypeKind::ARRAY:
                    localResult = applyArray(rows, args, context);
                    break;

                case TypeKind::MAP:
                    localResult = mapSubscript_.applyMap(rows, args, context);
                    break;

                default:
                    POLLUX_UNREACHABLE();
            }
            context.moveOrCopyResult(localResult, rows, result);
        }

        VectorPtr applyArray(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args,
            exec::EvalCtx &context) const {
            POLLUX_CHECK_EQ(args[0]->type_kind(), TypeKind::ARRAY);

            auto arrayArg = args[0];
            auto indexArg = args[1];

            switch (indexArg->type_kind()) {
                case TypeKind::INTEGER:
                    return applyArrayTyped<int32_t>(rows, arrayArg, indexArg, context);

                case TypeKind::BIGINT:
                    return applyArrayTyped<int64_t>(rows, arrayArg, indexArg, context);

                default:
                    POLLUX_UNSUPPORTED(
                        "Unsupported type for element_at index {}",
                        mapTypeKindToName(indexArg->type_kind()));
            }
        }

        static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
            std::vector<std::shared_ptr<exec::FunctionSignature> > signatures;

            // array(T), integer|bigint -> T
            for (const auto &indexType: {"integer", "bigint"}) {
                signatures.push_back(exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("T")
                    .argumentType("array(T)")
                    .argumentType(indexType)
                    .build());
            }

            // map(K,V), K -> V
            signatures.push_back(exec::FunctionSignatureBuilder()
                .typeVariable("K")
                .typeVariable("V")
                .returnType("V")
                .argumentType("map(K,V)")
                .argumentType("K")
                .build());

            return signatures;
        }

        template<typename I>
        VectorPtr applyArrayTyped(
            const SelectivityVector &rows,
            const VectorPtr &arrayArg,
            const VectorPtr &indexArg,
            exec::EvalCtx &context) const {
            auto *pool = context.pool();

            BufferPtr indices = allocate_indices(rows.end(), pool);
            auto rawIndices = indices->asMutable<vector_size_t>();

            // Create nulls for lazy initialization.
            NullsBuilder nullsBuilder(rows.end(), pool);

            exec::LocalDecodedVector arrayHolder(context, *arrayArg, rows);
            auto decodedArray = arrayHolder.get();
            auto baseArray = decodedArray->base()->as<ArrayVector>();
            auto arrayIndices = decodedArray->indices();

            exec::LocalDecodedVector indexHolder(context, *indexArg, rows);
            auto decodedIndices = indexHolder.get();

            auto rawSizes = baseArray->rawSizes();
            auto rawOffsets = baseArray->rawOffsets();

            // Optimize for constant encoding case.
            if (decodedIndices->isConstantMapping()) {
                bool allFailed = false;
                // If index is invalid, capture the error and mark all rows as failed.
                bool isZeroSubscriptError = false;
                const auto adjustedIndex =
                        adjustIndex(decodedIndices->value_at<I>(0), isZeroSubscriptError);
                if (isZeroSubscriptError) {
                    context.setErrors(rows, detail::zeroSubscriptError());
                    allFailed = true;
                }

                if (!allFailed) {
                    rows.applyToSelected([&](auto row) {
                        const auto elementIndex = getIndex(
                            adjustedIndex, row, rawSizes, rawOffsets, arrayIndices, context);
                        rawIndices[row] = elementIndex;
                        if (elementIndex == -1) {
                            nullsBuilder.set_null(row);
                        }
                    });
                }
            } else {
                rows.applyToSelected([&](auto row) {
                    const auto originalIndex = decodedIndices->value_at<I>(row);
                    bool isZeroSubscriptError = false;
                    const auto adjustedIndex =
                            adjustIndex(originalIndex, isZeroSubscriptError);
                    if (isZeroSubscriptError) {
                        context.setPolluxExceptionError(row, detail::zeroSubscriptError());
                        return;
                    }
                    const auto elementIndex = getIndex(
                        adjustedIndex, row, rawSizes, rawOffsets, arrayIndices, context);
                    rawIndices[row] = elementIndex;
                    if (elementIndex == -1) {
                        nullsBuilder.set_null(row);
                    }
                });
            }

            // Subscript into empty arrays always returns NULLs. Check added at the end
            // to ensure user error checks for indices are not skipped.
            if (baseArray->elements()->size() == 0) {
                return BaseVector::create_null_constant(
                    baseArray->elements()->type(), rows.end(), context.pool());
            }

            // Subscript can pass along very large elements vectors that can hold onto
            // memory and copy operations on them can further put memory pressure. We
            // try to flatten them if the dictionary layer is much smaller than the
            // elements vector.
            return BaseVector::wrap_in_dictionary(
                nullsBuilder.build(),
                indices,
                rows.end(),
                baseArray->elements(),
                true /*flattenIfRedundant*/);
        }

        // Normalize indices from 1 or 0-based into always 0-based (according to
        // indexStartsAtOne template parameter - no-op if it's false).
        template<typename I>
        vector_size_t adjustIndex(I index, bool &isZeroSubscriptError) const {
            // If array indices start at 1.
            if constexpr (indexStartsAtOne) {
                if (UNLIKELY(index == 0)) {
                    isZeroSubscriptError = true;
                    return 0;
                }

                // If larger than zero, adjust it.
                if (index > 0) {
                    index--;
                }
            }
            return index;
        }

        // Returns the actual Vector index given an array index. Checks and adjusts
        // negative indices, in addition to bound checks.
        // `index` is always a 0-based array index (see `adjustIndex` function
        // above).
        template<typename I>
        vector_size_t getIndex(
            I index,
            vector_size_t row,
            const vector_size_t *rawSizes,
            const vector_size_t *rawOffsets,
            const vector_size_t *indices,
            exec::EvalCtx &context) const {
            auto arraySize = rawSizes[indices[row]];

            if (index < 0) {
                // Check if we allow negative indices. If so, adjust.
                if constexpr (allowNegativeIndices) {
                    if constexpr (nullOnNegativeIndices) {
                        return -1;
                    } else {
                        index += arraySize;
                    }
                } else {
                    context.setPolluxExceptionError(row, detail::negativeSubscriptError());
                    return -1;
                }
            }

            // Check if index is within bound.
            if ((index >= arraySize) || (index < 0)) {
                // If we allow it, return null.
                if constexpr (allowOutOfBound) {
                    return -1;
                } else {
                    context.setPolluxExceptionError(row, detail::badSubscriptError());
                    return -1;
                }
            }

            // Resultant index is the sum of the offset in the input array and the
            // index.
            return rawOffsets[indices[row]] + index;
        }

    private:
        detail::MapSubscript mapSubscript_;
    };
} // namespace kumo::pollux::functions
